from fastapi.exceptions import HTTPException
import grpc
from functools import wraps
from loguru import logger
import asyncio
from typing import Callable, Any
from utils.status_code import get_http_code

def grpc_error_handler(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            return await func(*args, **kwargs)
        except grpc.RpcError as e:
            raise HTTPException(status_code=get_http_code(e.code()), detail=e.details())
        except Exception as e:
            logger.exception(e)
            raise HTTPException(status_code=500, detail=str(e))
    return wrapper

def grpc_retry_handler(max_retries: int = 2, retry_delay: float = 1.0):
    """
    gRPC服务调用重试装饰器，支持服务不可用时的自动重试和服务重新发现
    
    Args:
        max_retries: 最大重试次数，默认3次
        retry_delay: 重试延迟时间（秒），默认1秒
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs) -> Any:
            last_exception = None
            
            for attempt in range(max_retries + 1):
                try:
                    return await func(*args, **kwargs)
                    
                except grpc.RpcError as e:
                    last_exception = e
                    
                    # 检查是否是服务不可用的错误
                    if e.code() in [grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.DEADLINE_EXCEEDED]:
                        logger.warning(f"gRPC服务不可用 (尝试 {attempt + 1}/{max_retries + 1}): {e.details()}")
                        
                        if attempt < max_retries:
                            # 触发服务重新发现
                            await _refresh_service_discovery()
                            
                            # 等待一段时间后重试
                            await asyncio.sleep(retry_delay * (attempt + 1))  # 递增延迟
                            logger.info(f"准备进行第 {attempt + 2} 次重试...")
                            continue
                        else:
                            logger.error(f"达到最大重试次数 {max_retries}，服务调用失败")
                    else:
                        # 对于其他类型的gRPC错误，不进行重试
                        logger.error(f"gRPC调用失败，错误类型: {e.code()}, 详情: {e.details()}")
                        break
                        
                except Exception as e:
                    last_exception = e
                    logger.exception(f"服务调用异常 (尝试 {attempt + 1}/{max_retries + 1}): {e}")
                    
                    if attempt < max_retries:
                        await asyncio.sleep(retry_delay)
                        continue
                    else:
                        break
            
            # 所有重试都失败了，抛出最后的异常
            if isinstance(last_exception, grpc.RpcError):
                raise HTTPException(
                    status_code=get_http_code(last_exception.code()), 
                    detail=f"服务调用失败，已重试 {max_retries} 次: {last_exception.details()}"
                )
            else:
                logger.exception(last_exception)
                raise HTTPException(
                    status_code=500, 
                    detail=f"服务调用异常，已重试 {max_retries} 次: {str(last_exception)}"
                )
                
        return wrapper
    return decorator

async def _refresh_service_discovery():
    """
    刷新服务发现，重新获取可用的服务实例
    """
    try:
        from utils.tll_consul import ConsulClient
        consul_client = ConsulClient()
        
        logger.info("开始刷新服务发现...")
        
        # 重新获取健康的服务实例
        addresses = await consul_client.fetch_user_service_addresses(only_healthy=True)
        
        if addresses:
            logger.info(f"服务发现刷新成功，发现 {len(addresses)} 个健康服务实例")
        else:
            logger.warning("服务发现刷新后仍未找到健康的服务实例")
            
    except Exception as e:
        logger.error(f"刷新服务发现时发生异常: {e}")